from search.utils import *
from search.algorithm import *

class man_min(Heuristic):
    
    def __repr__(self):
        return 'man min' 
    
    def evaluate(self, problem_state):
        if problem_state.isGoal():return 0
        min = infinity
        for robot in problem_state.robots:
            x1, y1 = robot
            for dirt in problem_state.dirt_locations:
                x,y = dirt
                temp = abs(x-x1)+abs(y-y1)
                if temp < min: min = temp
        return min
    

class man_max(Heuristic):
    
    def __repr__(self):
        return 'man max' 
    
    def evaluate(self, problem_state):
        if problem_state.isGoal():return 0
        max = 0
        for robot in problem_state.robots:
            x1, y1 = robot
            for dirt in problem_state.dirt_locations:
                x,y = dirt
                temp = abs(x-x1)+abs(y-y1)
                if temp > max: max = temp
        return max    

class fast_manhattan_dirt(Heuristic):
    '''
    each dirt selects the closest robot and calculates the Manhattan distance to it.
    after that the robot is removed, and the next dirt selects the closet robot and calculates the Manhattan distacne to it...
    returning the sum of all of the Manhattan distances.
    '''
    def __repr__(self):
        return 'fst man drt'
        
    def evaluate(self, problem_state):
        def distance(a,b):
            return abs(a[0]-b[0])+abs(a[1]-b[1])
        sum = 0
        dirts = list(problem_state.dirt_locations)
        while(dirts):
            for robot in problem_state.robots:
                dis = infinity
                if not dirts: break
                for dirt in dirts:
                    temp = distance(dirt, robot)
                    if temp < dis: 
                        dis = temp
                        dirt_to_remove = dirt
                
                if not dirts: break
                dirts.remove(dirt_to_remove)
                sum += dis
        return sum 
            
class manhattan_dirt(Heuristic):
    '''
    each dirt selects the closest robot and calculates the Manhattan distance to it.
    returning the sum of all of the Manhattan distances.
    '''
    def __repr__(self):
        return 'man drt'
         
    def evaluate(self, problem_state):
        def distance(dirt, robots):
            distance = infinity
            x,y = dirt
            for robot in robots:
                x1, y1 = robot
                temp = abs(x-x1)+abs(y-y1)
                if temp < distance: distance = temp
            return distance
        sum = 0
        for dirt in problem_state.dirt_locations:
            sum += distance(dirt, problem_state.robots)
        return sum  

class scaled_manhattan_dirt(Heuristic):
    '''
    like manhattan_dirt but factored by a p
    '''
    def __init__(self, p):
        self.p = p
    
    def __repr__(self):
        return 'scld man drt, p=' + str(self.p)
    
    def evaluate(self, problem_state):
        def distance(dirt, robots):
            distance = infinity
            x,y = dirt
            for robot in robots:
                x1, y1 = robot
                temp = abs(x-x1)+abs(y-y1)
                if temp < distance: distance = temp
            return distance
        sum = 0
        for dirt in problem_state.dirt_locations:
            sum += distance(dirt, problem_state.robots)
        return self.p*sum
    
class averaged_scaled_manhattan_dirt(Heuristic):
    '''
    like scaled_manhattan_dirt but averaged by the number of dirts
    '''
    def __init__(self, p):
        self.p = p
    
    def __repr__(self):
        return 'avg scld man drt, p=' + str(self.p)
        
    def evaluate(self, problem_state):
        def distance(dirt, robots):
            distance = infinity
            x,y = dirt
            for robot in robots:
                x1, y1 = robot
                temp = abs(x-x1)+abs(y-y1)
                if temp < distance: distance = temp
            return distance
        sum = 0
        for dirt in problem_state.dirt_locations:
            sum += distance(dirt, problem_state.robots)
        if len(problem_state.dirt_locations) == 0: return 0 
        else: return float(self.p)*sum/len(problem_state.dirt_locations)  

class h(Heuristic):
    def evaluate(self, problem_state):
        return 1